import json

from server.module_chat.chat.core.agent.tool.api.chat_api import ChatApi
from server.module_chat.chat.core.agent.tool.api.chat_result import ChatResult
from server.module_chat.chat.core.agent.tool.tool import ToolApi


class DifyChatApi(ChatApi):

    def __init__(self, tool_api: ToolApi, **kwargs):
        super().__init__(tool_api, **kwargs)

    def build_data(self, **kwargs) -> dict:
        data = {
            "inputs": {},
            "query": kwargs["query"],
            "conversation_id": kwargs["conversation_id"],
            "user": "test",
            "response_mode": "streaming" if kwargs["stream"] else "blocking",
        }
        if kwargs["conversation_id"] is not None:
            data["conversation_id"] = kwargs["conversation_id"]
        return data

    def build_result(self, chunk: bytes) -> ChatResult | None:
        chunk_str = chunk.decode("utf-8")
        # 判断chunk_str是否为空
        if chunk_str:
            try:
                result = json.loads(chunk_str.lstrip('data:').strip())
                if result["event"] == "message":
                    return ChatResult(**{
                        "conversation_id": result["conversation_id"],
                        "message_id": result["message_id"],
                        "text": result["answer"]
                    })
            except Exception as e:
                print(f"Error parsing chunk({chunk_str}): {e}")
        return None
